import tushare as ts
import pymysql
import pandas as pd
import concurrent.futures
import time
from datetime import datetime, timedelta


class DataFetcher:
    def __init__(self, token):
        self.token = token

    def fetch_stock_data(self, stock, start_date, end_date, data_type='daily'):
        # 初始化 Tushare API
        pro = ts.pro_api()
        time.sleep(0.075)
        try:
            stock_data = pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
            return stock_data
        except Exception as e:
            print(e)

    def fetch_data(self, symbol, start_date, end_date, data_type='daily') -> dict:
        # 初始化 Tushare API
        ts.set_token(self.token)
        pro = ts.pro_api()
        data = dict()
        all_stocks_data = pd.DataFrame()

        if data_type == 'daily':
            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
                # 提交每个股票的任务给线程池
                futures = [executor.submit(self.fetch_stock_data, stock, start_date, end_date, data_type)
                           for stock in symbol]

                # 获取已完成的任务结果
                for future in concurrent.futures.as_completed(futures):
                    stock_data = future.result()
                    all_stocks_data = pd.concat([all_stocks_data, stock_data])

            data['data'] = all_stocks_data
            data['type'] = 'daily'
        else:
            raise ValueError(f"Invalid data_type: {data_type}")

        return data

    def fetch_turnover_data(self, stock, trade_date):
        pro = ts.pro_api()
        time.sleep(0.3)
        try:
            turnover_data = pro.daily_basic(ts_code=stock, trade_date=trade_date,
                                            fields='ts_code, trade_date, turnover_rate')
            return turnover_data
        except Exception as e:
            print(f"Error occurred: {e}")

    def fetch_turnover(self, symbol, start_date_string, end_date_string) -> dict:
        ts.set_token(self.token)
        data = dict()
        all_turnover_data = pd.DataFrame()
        tasks = []
        start_date = datetime.strptime(start_date_string, '%Y%m%d')
        end_date = datetime.strptime(end_date_string, '%Y%m%d')
        for stock in symbol:
            current_date = start_date
            while current_date <= end_date:
                current_date_string = current_date.strftime('%Y%m%d')
                tasks.append((stock, current_date_string))
                current_date += timedelta(days=1)  # 可能存在是周末的问题

        print('Download', len(tasks), ' data')
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
            # 提交每个票的任务给线程池
            futures = [executor.submit(self.fetch_turnover_data, stock, trade_date)
                       for stock, trade_date in tasks]

            # 获取已完成的任务结果
            for future in concurrent.futures.as_completed(futures):
                if future.result() is not None:
                    turnover_data = future.result()
                    all_turnover_data = pd.concat([all_turnover_data, turnover_data])
            data['data'] = all_turnover_data
        return data
